# -*- coding: utf-8 -*-
# -------------------------------------------------
# File Name： tsp
# Description :
# Author : 'li'
# date： 2022/5/31
# -------------------------------------------------
# Change Activity:
# 2022/5/31:
# -------------------------------------------------
"""
This file is used to simulate TSP environment.
"""
import torch

from ml.rl.envs.tsp.tsp_dataset import TSPDataset


class TSP:
    """
    Provides an environment for TSP
    """

    def __init__(self, points_size=50, num_samples=10000):
        """

        Args:
            points_size: the number of points per env graph.
            num_samples:  total number of training samples.
        """
        self.points_size = points_size
        self.num_samples = num_samples

    def create_tsp_dataset(self):
        return TSPDataset(points_size=self.points_size, num_samples=self.num_samples)

    @staticmethod
    def get_costs(dataset, pi):
        # Check that tours are valid, i.e. contain 0 to n -1
        assert (
                torch.arange(pi.size(1), out=pi.data.new()).view(1, -1).expand_as(pi) ==
                pi.data.sort(1)[0]
        ).all(), "Invalid tour"

        # Gather dataset in order of tour
        d = dataset.gather(1, pi.unsqueeze(-1).expand_as(dataset))

        # Length is distance (L2-norm of difference) from each next location from its prev and of last from first
        return (d[:, 1:] - d[:, :-1]).norm(p=2, dim=2).sum(1) + (d[:, 0] - d[:, -1]).norm(p=2, dim=1), None
